package com.soushangxing.comfyuibackendv2.service

import com.soushangxing.comfyuibackendv2.bean.GlobalQueue
import com.soushangxing.comfyuibackendv2.bean.ResponseType
import com.soushangxing.comfyuibackendv2.entity.OrderEntity
import com.soushangxing.comfyuibackendv2.repository.OrderErrorRepository
import com.soushangxing.comfyuibackendv2.repository.OrderRepository
import kotlinx.coroutines.delay
import org.springframework.beans.factory.annotation.Autowired
import org.springframework.beans.factory.annotation.Value
import org.springframework.stereotype.Service
import org.springframework.web.client.RestTemplate
import org.springframework.web.multipart.MultipartFile
import org.springframework.core.io.ByteArrayResource
import org.springframework.http.*
import org.springframework.util.LinkedMultiValueMap
import java.io.OutputStream
import java.net.HttpURLConnection
import java.net.URL
import kotlin.math.abs
import kotlin.random.Random

@Service
class ComfyUIService {

    @Autowired
    lateinit var orderRepository: OrderRepository
    @Autowired
    lateinit var orderErrorRepository: OrderErrorRepository

    @Autowired
    lateinit var fileService: FileService

    @Autowired
    lateinit var globalQueue: GlobalQueue

    /**
     * 通过阿里云上传文件
     */
    fun uploadFileAliyun(file:MultipartFile, type:String = "input") : ResponseType{
        return ResponseType(fileService.saveFile(file,type),0,"上传完成")
    }

    /**
     * 向服务器下发一个任务【新版】
     * 仅下发，不等待和轮询检查结果
     */
    fun queuePromptNew(taskId:String, userId:String, prompt:String,
                       serverUrl:String, clientId:String) : ResponseType{
        try{
            val random = Random(System.currentTimeMillis())
            val randomInt = random.nextInt()

            val url = URL(serverUrl+"/prompt")
            // 将随机数、client替换
            val finalPrompt = prompt.replace("*@#seed#@*",abs(randomInt).toString())
                .replace("*@#client_id#@*",clientId)

            val jsonBody = finalPrompt

            with(url.openConnection() as HttpURLConnection) {
                requestMethod = "POST"
                doOutput = true
                setRequestProperty("Content-Type", "application/json")
                setRequestProperty("charset", "utf-8")
                setRequestProperty("Content-Length", jsonBody.toByteArray().size.toString())

                val outputStream: OutputStream = outputStream
                outputStream.write(jsonBody.toByteArray())
                outputStream.flush()
                outputStream.close()

                if (responseCode == HttpURLConnection.HTTP_OK) {
                    inputStream.bufferedReader().use {
                        val response = it.readText()
                        val pattern = """("prompt_id":\s*"([^"]+)")""".toRegex()
                        val matchResult = pattern.find(response)
                        val promptId = matchResult?.groupValues?.get(2)
                        if (promptId.isNullOrEmpty()){
                            return ResponseType("请求ID没有成功获取到，请求失败")
                        } else {
                            println("TASK START: promptId: $promptId at server :$serverUrl ")
                            val orderEntity = orderRepository.findFirstByPromptId(promptId)?:OrderEntity()
                            orderEntity.promptId = promptId
                            orderEntity.taskId = taskId
                            orderEntity.status = "进行中"
                            orderEntity.userId = userId
                            orderEntity.serverUrl = serverUrl
                            orderRepository.save(orderEntity)

                            val task = globalQueue.globalBackendQueue.find { it.taskId == taskId }
                            if (task != null) {
                                task.promptId = promptId
                                task.status = "进行中"
                            }

                            return ResponseType(orderEntity,0,"任务下达成功，请用任务ID查询进度")

                        }
                    }
                } else {
                    println("TASK FAILED: Failed to post data")
                }
            }
            return ResponseType(null,1,"运行失败，没有获得结果。")

        } catch (e:Exception){
            return ResponseType("运行失败:"+e.message)
        }
    }


    /**
     * 入口：根据type参数分配指定的处理服务。
     * 由于每个服务的入参、替换内容可能不同，以本函数为校验和调度。
     */
    fun getProgress(taskId:String, userId:String): ResponseType{
        // 先从全局队列中找此任务，找得到就优先返回全局队列中的
        val task = globalQueue.globalBackendQueue.find { it.taskId == taskId }
        if (task != null) {
            if (task.userId != userId) {
                return ResponseType("不可以查看他人的任务进度")
            }
            return ResponseType(task,0,"任务查询完毕")
        }

        // 如果全局队列中找不到本任务，再尝试从数据库中查询
        val order = orderRepository.findFirstByTaskId(taskId)
        if (order != null){
            if (order.userId != userId) {
                return ResponseType("不可以查看他人的任务进度")
            }
            if (order.output.isNotEmpty() && order.outputPath.isEmpty()){
                viewImage(order.output,order.taskId)
            }
            return ResponseType(order,0,"任务查询完毕")
        }

        // 如果数据库的已完成任务里也没有，尝试在错误任务堆里找
        val error = orderErrorRepository.findFirstByTaskId(taskId) ?: return ResponseType("未找到此任务")
        if (error.userId != userId) {
            return ResponseType("不可以查看他人的任务进度")
        }
        return ResponseType(error,0,"任务查询完毕")

    }



    /**
     * 查看图片（从服务器请求）
     * 并将图片存到阿里云，形成外链返回
     */
    fun viewImage(filename:String, taskId: String):ResponseType{
        try{
            val order = orderRepository.findFirstByTaskId(taskId)
            if (order == null) {
                return ResponseType("任务不存在，无法查询")
            }
            val serverUrl = order.serverUrl
            val restTemplate = RestTemplate()
            val url = serverUrl+"/view?filename=$filename&type=output"
            val headers = HttpHeaders()

            // 创建请求实体
            val requestEntity: HttpEntity<ByteArray> = HttpEntity(null, headers)

            // 发送请求并接收响应
            val responseEntity: ResponseEntity<ByteArray> = restTemplate.exchange(url, HttpMethod.GET, requestEntity, ByteArray::class.java)

            // 获取响应体
            val imageBytes = responseEntity.body
            if (imageBytes == null){
                return ResponseType("未返回图片")
            }
            val file = fileService.saveFile(imageBytes,"output")
            if (file != null) {
                println("TASK SAVED: Result filename: ${filename}, Saving to ${file.path}")
                if (order != null) {
                    order.output = filename
                    order.status = "完成"
                    order.outputPath = file.path
                    order.previewPath = file.path+"?x-oss-process=image/resize,m_fill,w_{width},h_{height}"
                    orderRepository.save(order)
                }
            } else {
                return ResponseType("图片存储失败")
            }

            return ResponseType(file, 0, "图片生成完成")
        }catch (e:Exception){
            return ResponseType("预览图片失败:"+e.message)
        }
    }

}

